'''
MMDetection目标检测
in: str

path to image
---
out: list[80] -> np.ndarray(num,5)

预训练在COCO数据集，共80类
每一类检测出num个bbox
每个bbox有5个字段 (左上点x，左上点y，右下点x，右下点y，置信概率）
'''

from mmdet.apis import init_detector, inference_detector
import torch


def get_label_list(self):
    CLASSES = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
               'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
               'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
               'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
               'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
               'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
               'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
               'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
               'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
               'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
               'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
               'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
               'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
               'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush')
    return CLASSES


class MMDetection(torch.nn.Module):
    def __init__(self, device='cpu'):
        super().__init__()
        self.device = device
        self.config_file = 'utils/MMDetection/configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py'
        self.checkpoint_file = 'utils/MMDetection_checkpoint_file/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth'
        self.model = init_detector(self.config_file, self.checkpoint_file, self.device)

    def forward(self, img):
        out = inference_detector(self.model, img)

        return out
